import asyncio
import json
from typing import Any
from fastapi import APIRouter, FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI(root_path='/api/v1')

origins = [
    "http://localhost:5173"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

chatApi = APIRouter(prefix='/chat')

class Message(BaseModel):
    id: str
    content: str
    role: str = 'user'

@chatApi.post('/completion')
async def completion(msg: Message) -> Any:
    print(msg)
    
    def stream_reply(data: dict) -> str:
        return f'data: {json.dumps(data)}\n\n'

    async def stream():
        content = msg.content
        tmp = ""
        for c in content:
            tmp += c
            await asyncio.sleep(0.2)
            yield stream_reply({"data": tmp})
        yield stream_reply({"data": True})

    return StreamingResponse(
        stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive"
        }
    )

app.include_router(chatApi)
